#ifdef __cplusplus
extern "C" {
#endif

#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <string.h>
#include <pthread.h>

#include "thread.h"

#define PORT 8080
#define BUFFER_SIZE 1024

extern char Server_file_path[64];  // Path for the server to save the received file
extern char Client_file_path[64];  // Path for the client to send the file

// TCP Server function
int tcp_server()
{
    printf("TCP 服务器已启动\n");

    int server_fd, connect_fd, file_fd;
    struct sockaddr_in address, client_addr;
    char buffer[BUFFER_SIZE];
    int bytes_received;

    // 1. Create server socket
    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) == 0) {
        perror("创建套接字时发生错误");
        exit(EXIT_FAILURE);
    }

    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(PORT);

    // 2. Bind the socket
    if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("绑定套接字时失败");
        exit(EXIT_FAILURE);
    }

    // 3. Listen for incoming client connections
    if (listen(server_fd, 3) < 0) {
        perror("监听连接时发生错误");
        exit(EXIT_FAILURE);
    }

    while (1) {
        // 4. Accept a client connection
        if ((connect_fd = accept(server_fd, (struct sockaddr *)&client_addr, (socklen_t*)&client_addr)) < 0) {
            perror("接受客户端连接失败");
            exit(EXIT_FAILURE);
        }

        printf("客户端已连接，正在等待文件...\n");

        // Open file to save received data
        file_fd = open(Server_file_path, O_WRONLY | O_CREAT, 0644);
        if (file_fd == -1) {
            perror("打开文件进行写入时发生错误");
            exit(EXIT_FAILURE);
        }

        // Receive file data from client
        while ((bytes_received = recv(connect_fd, buffer, BUFFER_SIZE, 0)) > 0) {
            if (write(file_fd, buffer, bytes_received) != bytes_received) {
                perror("写入接收到的数据时发生错误");
            }
        }

        // Close the connection and file when done
        if (bytes_received == 0) {
            printf("客户端已断开连接。\n");
            close(connect_fd);
            close(file_fd);
        } else {
            perror("接收数据时发生错误");
            close(file_fd);
        }
    }

    close(server_fd);
    return 0;
}

// TCP Client function
int tcp_client()
{
    sleep(1);  // Allow time for the server to start
    printf("TCP 客户端已启动\n");

    int socket_fd, file_fd;
    struct sockaddr_in server_addr;
    char buffer[BUFFER_SIZE];
    int bytes_read;

    // 1. Create client socket
    if ((socket_fd = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
        perror("创建套接字时发生错误");
        exit(EXIT_FAILURE);
    }

    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    if (inet_pton(AF_INET, "127.0.0.1", &server_addr.sin_addr) <= 0) {
        perror("提供的地址无效");
        exit(EXIT_FAILURE);
    }

    // 2. Establish connection to the server
    if (connect(socket_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("连接到服务器失败");
        exit(EXIT_FAILURE);
    }

    printf("已连接到服务器，开始传输文件...\n");

    // Open the file to send
    file_fd = open(Client_file_path, O_RDONLY);
    if (file_fd == -1) {
        perror("打开文件以进行发送时发生错误");
        exit(EXIT_FAILURE);
    }

    // Send file data to the server
    while ((bytes_read = read(file_fd, buffer, BUFFER_SIZE)) > 0) {
        if (send(socket_fd, buffer, bytes_read, 0) != bytes_read) {
            perror("文件传输过程中发生错误");
            exit(EXIT_FAILURE);
        }
    }

    // Close the file and socket
    close(file_fd);
    close(socket_fd);
    pthread_exit(NULL);

    return 0;
}

// UDP Server function
int udp_server()
{
    printf("UDP 服务器已启动\n");

    int server_fd, file_fd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t addr_len = sizeof(client_addr);
    char buffer[BUFFER_SIZE];
    int bytes_received;

    // 1. Create server socket
    if ((server_fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("创建套接字时发生错误");
        exit(EXIT_FAILURE);
    }

    // 2. Set up server address
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(PORT);

    // 3. Bind the socket
    if (bind(server_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("绑定套接字时失败");
        exit(EXIT_FAILURE);
    }

    // Open file to save received data
    file_fd = open(Server_file_path, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (file_fd == -1) {
        perror("打开文件进行写入时发生错误");
        exit(EXIT_FAILURE);
    }

    while (1) {
        // Receive data from the client
        bytes_received = recvfrom(server_fd, buffer, BUFFER_SIZE, 0, (struct sockaddr *)&client_addr, &addr_len);
        if (bytes_received < 0) {
            perror("接收数据时发生错误");
            break;
        }

        // Write the received data to the file
        if (write(file_fd, buffer, bytes_received) != bytes_received) {
            perror("写入文件时发生错误");
            break;
        }

        printf("已接收到来自客户端的数据。\n");

        // Check for the "END" message to indicate the end of the transfer
        if (bytes_received == 3 && strncmp(buffer, "END", 3) == 0) {
            printf("文件传输完成。\n");
            break;
        }
    }

    close(file_fd);
    close(server_fd);
    pthread_exit(NULL);

    return 0;
}

// UDP Client function
int udp_client()
{
    sleep(1);  // Allow time for the server to start
    printf("UDP 客户端已启动\n");

    int socket_fd, file_fd;
    struct sockaddr_in server_addr;
    char buffer[BUFFER_SIZE];
    int bytes_read;

    // 1. Create client socket
    if ((socket_fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("创建套接字时发生错误");
        exit(EXIT_FAILURE);
    }

    // 2. Set up server address
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    if (inet_pton(AF_INET, "127.0.0.1", &server_addr.sin_addr) <= 0) {
        perror("提供的服务器地址无效");
        exit(EXIT_FAILURE);
    }

    // 3. Open the file to send
    file_fd = open(Client_file_path, O_RDONLY);
    if (file_fd == -1) {
        perror("打开文件以进行发送时发生错误");
        exit(EXIT_FAILURE);
    }

    printf("正在向服务器发送文件数据...\n");

    // Send the file data to the server
    while ((bytes_read = read(file_fd, buffer, BUFFER_SIZE)) > 0) {
        if (sendto(socket_fd, buffer, bytes_read, 0, (struct sockaddr *)&server_addr, sizeof(server_addr)) != bytes_read) {
            perror("发送数据到服务器时发生错误");
            exit(EXIT_FAILURE);
        }
    }

    // Send "END" message to mark the end of file transfer
    const char *end_message = "END";
    sendto(socket_fd, end_message, strlen(end_message), 0, (struct sockaddr *)&server_addr, sizeof(server_addr));

    printf("文件传输已完成。\n");

    // Close file and socket
    close(file_fd);
    close(socket_fd);
    pthread_exit(NULL);

    return 0;
}

#ifdef __cplusplus
}
#endif
